from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage
import ast
import json
import sys
import pymysql
import os

# 星火认知大模型调用秘钥信息
SPARKAI_URL = 'wss://spark-api.xf-yun.com/v3.1/chat'
SPARKAI_APP_ID = '1d38b17a'
SPARKAI_API_SECRET = 'MTJjOGQ5YzE1ZWZjYjJkOWUxZTM4Mzcx'
SPARKAI_API_KEY = 'b33d686049c006fd12a7e9b345ec0409'
SPARKAI_DOMAIN = 'generalv3'

os.environ["http_proxy"] = "http://127.0.0.1:10808"
os.environ["https_proxy"] = "http://127.0.0.1:10808"

def dataBase(questionID):
    connection = pymysql.connect(
        host="www.yym-free.com",
        user="fuchuang",
        password="fuchuang",
        db="fuchuang",
        port=3306,
    )
    cur = connection.cursor()
    cur.execute("SELECT * FROM dialog WHERE questionID = %s", (questionID,))
    all = cur.fetchall()
    new_list = []
    for i in all:
        new_item = {
            "askStr": i[2],
            "answerStr": i[3],
        }
        new_list.append(new_item)
        print(new_item)
    connection.commit()
    connection.close()
    return new_list

def list_to_chat_history(chat_history: list) -> list:
    from langchain.memory import ChatMessageHistory
    message_history = ChatMessageHistory()
    for entry in chat_history:
        message_history.add_user_message(entry["askStr"])
        message_history.add_ai_message(entry["answerStr"])
    return message_history

def convert_chat_history_to_list(message_history: ChatMessageHistory) -> list:
    chat_history_list = []
    for message in message_history.messages:
        if isinstance(message, HumanMessage):
            chat_history_list.append({"sender": "User", "message": message.content})
        elif isinstance(message, AIMessage):
            chat_history_list.append({"sender": "LLM", "message": message.content})
    return chat_history_list

def chat_with_llm(user_input: str, chat_history: list) -> dict:
    try:
        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", "You are a helpful assistant. Answer all questions to the best of your ability."),
                MessagesPlaceholder(variable_name="messages"),
            ]
        )
        llm = ChatSparkLLM(
            spark_api_url=SPARKAI_URL,
            spark_app_id=SPARKAI_APP_ID,
            spark_api_key=SPARKAI_API_KEY,
            spark_api_secret=SPARKAI_API_SECRET,
            spark_llm_domain=SPARKAI_DOMAIN,
            streaming=False,
        )
        output_parser = StrOutputParser()
        chain = prompt | llm | output_parser
        message_history = list_to_chat_history(chat_history)
        message_history.add_user_message(user_input)
        response = chain.invoke(
            {
                "messages": message_history.messages,
            }
        )
        chat_history.append({"askStr": user_input, "answerStr": response})
        print(response)
    except Exception as e:
        print(e)

if __name__ == "__main__":
    user_input = sys.argv[1]
    questionID = sys.argv[2]
    quesID = int(questionID)
    list = dataBase(quesID)
    chat_with_llm(user_input, list)
